package com.disv.support.encrypt;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSONObject;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.openssl.PEMParser;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import java.io.IOException;
import java.io.StringReader;
import java.nio.charset.StandardCharsets;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Base64;

/**
 * 公钥：用于对固定密钥明文加密、验签
 *
 * @author caojian.cn@gmail.com
 * @date 2021/7/1
 */
@Slf4j
public class CertificatePublicKey {
    final String ENCODING = "UTF-8";
    protected PublicKey publicKey;
    protected String sigAlgName;

    /**
     * 增加一个返回算法名的判断
     *
     * @return 算法名称
     */
    private String getSignAlgName() {
        return "RSA".equals(this.publicKey.getAlgorithm()) ? "SHA256WITHRSA" : "SM3WITHSM2";
    }

    public CertificatePublicKey(@NonNull String certificateString) throws CertificateException {
        try (StringReader reader = new StringReader(certificateString);
             PEMParser pemParser = new PEMParser(reader)) {
            Object object = pemParser.readObject();
            if (object instanceof X509CertificateHolder) {
                X509Certificate certificate = new JcaX509CertificateConverter().setProvider("BC").getCertificate((X509CertificateHolder) object);
                sigAlgName = certificate.getSigAlgName();
                publicKey = certificate.getPublicKey();
            } else {
                throw new CertificateException("不支持的证书类型，目前只支持X509");
            }
        } catch (IOException ex) {
            throw new CertificateException("加载证书失败", ex);
        }
    }


    /**
     * 验签
     *
     * @param plaintext
     * @param signature
     * @return
     * @throws Exception
     */
    public boolean verify(@NonNull String plaintext, @NonNull String signature) throws Exception {
        Signature engine = Signature.getInstance(getSignAlgName());
        engine.initVerify(this.publicKey);
        engine.update(plaintext.getBytes(ENCODING));
        return engine.verify(Base64.getDecoder().decode(signature));
    }

    /**
     * 加密
     *
     * @param plainText 待加密文本
     * @return
     * @throws Exception
     */
    public String encrypt(String plainText) throws Exception {
        // 获取加密后的Json
        JSONObject body = createEncryptJson(plainText);
        //返回加密字符串：key#ciphertext
        return body.getString("key") + '.' + body.getString("data");
    }


    /**
     * 加密
     *
     * @param data
     * @return
     * @throws Exception
     */
    public JSONObject creatRequestBody(String data) throws Exception {
        if (StrUtil.isBlank(data)) {
            return new JSONObject();
        }
        JSONObject body = createEncryptJson(data);
        log.info("组装的请求body报文:{}", body);
        return body;
    }

    /**
     * 加密
     *
     * @param data 代加密字符串
     * @description 抽取原来的公共方法, 返回key, encryptMethod, data的JSONObject
     */
    private JSONObject createEncryptJson(String data) throws Exception {
        //用于加密固定密钥的公钥
        String keyAlgName = getKeyAlgName();
        log.info("加密过程，非对称密钥算法：" + keyAlgName);
        Cipher keyCipher = Cipher.getInstance(keyAlgName, "BC");
        keyCipher.init(Cipher.ENCRYPT_MODE, publicKey);
        String keyGeneratorName = getKeyGeneratorName();
        log.info("加密过程，固定密钥算法:" + keyGeneratorName);
        KeyGenerator generator = KeyGenerator.getInstance(keyGeneratorName, "BC");
        generator.init(getKeySize());
        SecretKey key = generator.generateKey();
        //产生一个随机密钥
        //随机密钥的算法
        String algName = getMsgAlgName();
        log.info("加密过程，固定密钥算法参数:" + algName);
        Cipher cipher = Cipher.getInstance(algName, "BC");
        cipher.init(Cipher.ENCRYPT_MODE, key, new SecureRandom());
        //加密报文:
        byte[] encryptDataBytes = cipher.doFinal(data.getBytes(ENCODING));
        //取IV
        byte[] iv = cipher.getIV();
        byte[] keyBytes = key.getEncoded();

        //将IV和Key放在密文前，组成发送报文
        byte[] combinedPayload = new byte[iv.length + keyBytes.length];
        System.arraycopy(iv, 0, combinedPayload, 0, iv.length);
        System.arraycopy(keyBytes, 0, combinedPayload, iv.length, keyBytes.length);
        //固定密钥明文
        String solidKey = new String(combinedPayload, StandardCharsets.UTF_8);
        //对固定密钥进行加密
        byte[] encryptedPayload = keyCipher.doFinal(combinedPayload);

        String encryptKeyBase64 = Base64.getEncoder().encodeToString(encryptedPayload);
        String encryptDataBase64 = Base64.getEncoder().encodeToString(encryptDataBytes);

        JSONObject body = new JSONObject();
        body.put("key", encryptKeyBase64);
        body.put("encryptMethod", keyGeneratorName);
        body.put("data", encryptDataBase64);
        return body;
    }

    /**
     * 获取密钥类型
     *
     * @return
     */
    private String getKeyAlgName() {
        return "SM3WITHSM2".equals(sigAlgName) ? "SM2" : "RSA/NONE/PKCS1Padding";
    }

    /**
     * 固定密钥加密算法
     *
     * @return
     */
    private String getMsgAlgName() {
        return "SM3WITHSM2".equals(sigAlgName) ? "SM4/CBC/PKCS5Padding" : "AES/CBC/PKCS5Padding";
    }

    private String getKeyGeneratorName() {
        return "SM3WITHSM2".equals(sigAlgName) ? "SM4" : "AES";
    }

    private int getKeySize() {
        return "SM3WITHSM2".equals(sigAlgName) ? 128 : 256;
    }
}
